import os
import glob
import cv2
import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import sys

from pylab import mpl
#mpl.rcParams['font.sans-serif'] = ['SimHei']
mpl.rcParams['font.sans-serif'] = ['Microsoft YaHei']# 指定默认字体：解决plot不能显示中文问题
mpl.rcParams['axes.unicode_minus'] = False           # 解决保存图像是负号'-'显示为方块的问题

try:
    import xml.etree.CElementTree as ET
except:
    import xml.etree.ElementTree as ET

def bbox_from_xml(path = 'Annotations'):
    if os.path.exists('bbox.txt'):
        np.loadtxt('bbox.txt',delimiter=',',dtype=float)

    module_input_size = 768
    wh_list = []
    center_list = []
    xml_path = []
    if os.path.isdir(path):
        xml_path = glob.glob(f"{path}/*.xml")
    if os.path.isfile(path):
        name_list = np.loadtxt(path,dtype=str)
        xml_path = [f"{path}/{name}.xml" for name in name_list]
            
    for xmlFilePath in xml_path:
        tree = ET.parse(xmlFilePath)
        root = tree.getroot()
        ImageSize_w = int(root.find('size').find('width').text)
        ImageSize_h = int(root.find('size').find('height').text)
        if not ImageSize_w or not ImageSize_h:
            filename = root.find('filename').text
            image = cv2.imread(f"JPEGImages/{filename}")
            ImageSize_w = image.shape[1]
            ImageSize_h = image.shape[0]
        for bbox in root.iter('bndbox'):
            xmax = int(bbox.find('xmax').text)  #find 只能查找直系第一个元素
            xmin = int(bbox.find('xmin').text)
            ymax = int(bbox.find('ymax').text)
            ymin = int(bbox.find('ymin').text)
            w = (xmax - xmin)*module_input_size/ImageSize_w
            h = (ymax - ymin)*module_input_size/ImageSize_h
            center_x = (xmax + xmin)/2*module_input_size/ImageSize_w
            center_y = (ymax + ymin)/2*module_input_size/ImageSize_h
            if w and h:
                wh_list.append([w,h])
                center_list.append([center_x,center_y])
                
    wh_list,center_list = np.array(wh_list),np.array(center_list)
    bbox = np.concatenate((wh_list,center_list),axis=1)
    with open('bbox.txt','a+') as f:
        np.savetxt(f,bbox,fmt='%d,%d,%d,%d',delimiter=',')
        
    return bbox

def bbox_from_npy(path = ''):
    #top_xmin, top_ymin, top_xmax, top_ymax
    bbox = []
    res = np.load(path,allow_pickle=True)
    for b in res.item()['cls']:
        if not b[0] == -1:
            w = (b[2] - b[0])
            h = (b[3] - b[1])
            center_x = (b[2] + b[0])/2
            center_y = (b[3] - b[1])/2
            bbox.append([w,h,center_x,center_y])

def show_xml_bbox_size_distribute():
    area_list = []
    for xmlFilePath in glob.glob("Annotations/*.xml"):
        tree = ET.parse(xmlFilePath)
        root = tree.getroot()
        for bbox in root.iter('bndbox'):
           area = float((int(bbox.find('xmax').text) - int(bbox.find('xmin').text))*(int(bbox.find('ymax').text) - int(bbox.find('ymin').text)))
           area_list.append(area)
        
    #area_list = [(area - min(area_list))*100.0/(max(area_list) - min(area_list)) for area in area_list]
    #area_list = [(area - np.mean(area_list))/np.std(area_list) for area in area_list]
    area_list = sorted(area_list,reverse=True)[:int(len(area_list)*0.5)]
    data = pd.DataFrame({"area":area_list})
    #print(data)
    cats1 = pd.cut(data['area'].values, bins= np.linspace(min(area_list), max(area_list), num=50))
    #cats1 = pd.qcut(data['area'].values, 10) #数据均分
    #print(cats1)
    counts = cats1.value_counts()
    counts_df = pd.DataFrame(counts, columns=['counts'])
    counts_df['frequency'] = counts_df / counts_df['counts'].sum()
    #pinshu_df['频率%'] = pinshu_df['频率f'].map(lambda x: '%.2f%%' % (x * 100))
    #pinshu_df['累计频率f'] = pinshu_df['频率f'].cumsum()
    #pinshu_df['累计频率%'] = pinshu_df['累计频率f'].map(lambda x: '%.4f%%' % (x * 100))
    counts_df['counts'].plot(kind='bar')
    plt.show()
   
def show_predict_bbox_size_distribute():
    area_list = []
    bbox = np.load("predict_skip_up.npy",allow_pickle=True)
    for b in bbox.item()['cls']:
        area = (b[2] - b[0])*(b[3] - b[1])
        area_list.append(area)
    
    #area_list = [(area - min(area_list))*100.0/(max(area_list) - min(area_list)) for area in area_list]
    data = pd.DataFrame({"area":area_list})
    cats1 = pd.cut(data['area'].values, bins=np.linspace(min(area_list), max(area_list), num=50))
    counts = cats1.value_counts()
    counts_df = pd.DataFrame(counts, columns=['counts'])
    counts_df['frequency'] = counts_df / counts_df['counts'].sum()
    #pinshu_df['频率%'] = pinshu_df['频率f'].map(lambda x: '%.2f%%' % (x * 100))
    #pinshu_df['累计频率f'] = pinshu_df['频率f'].cumsum()
    #pinshu_df['累计频率%'] = pinshu_df['累计频率f'].map(lambda x: '%.4f%%' % (x * 100))
    counts_df['counts'].plot(kind='bar')
    plt.show()


def show_xml_bbox_space_distribute(bbox):
    wh_list,center_list = bbox[:,:2],bbox[:,2:]
    
    # bbox中心点分布
    bbox_df = pd.DataFrame({"center_x":center_list[:,0],"center_y":center_list[:,1],"w":wh_list[:,0],"h":wh_list[:,1]})
    bbox_df.plot(kind='scatter',x='center_x',y='center_y',title='bbox中心点分布')
    plt.show()
    # y方向上的bbox数量分布
    y_cut_df = pd.cut(bbox_df['center_y'].values,bins=np.linspace(bbox_df['center_y'].min(),bbox_df['center_y'].max(), num=60))
    y_counts_df = y_cut_df.value_counts()
    y_counts_df.plot(kind='bar',title='y方向上的bbox数量分布')
    plt.show()
    # x方向上的bbox数量分布
    x_cut_df = pd.cut(bbox_df['center_x'].values,bins=np.linspace(bbox_df['center_x'].min(),bbox_df['center_x'].max(), num=60))
    x_counts_df = x_cut_df.value_counts()
    x_counts_df.plot(kind='bar',title='x方向上的bbox数量分布')
    plt.show()
    # x方向上的bbox面积分布
    bbox_df['area'] = pow(bbox_df['w']*bbox_df['h'],0.25)
    bbox_df.plot(kind='scatter',x='center_y',y="area",title='y方向上的bbox面积分布')
    plt.show()
    # bbox面积分布
    area_cut_df = pd.cut(bbox_df['area'].values,bins=np.linspace(bbox_df['area'].min(),bbox_df['area'].max(), num=60))
    area_counts_df = area_cut_df.value_counts()
    area_counts_df.plot(kind='bar',title='bbox面积分布')
    plt.show()
    # bbox中心点和大小分布,
    #DataFrame.corr(method='pearson', min_periods=1)#分析数据的相关性，线性或非线性
    #heat_df = bbox_df.loc['center_x','center_y','area']
    #sns.heatmap(bbox_df['area'],)
    #"""将数值映射为颜色"""
    norm = mpl.colors.Normalize(vmin=np.min(bbox_df['area']), vmax=np.max(bbox_df['area']))
    cmap  = mpl.cm.get_cmap('rainbow')
    clors = [cmap(norm(val)) for val in bbox_df['area']]
    bbox_df.plot(kind='scatter',x='center_x',y='',s=bbox_df['area'],color=clors,title='bbox中心点和大小分布')
    plt.show()
    
#
bbox = []
if len(sys.argv) >= 2:
    if sys.argv[1].endswith('.npy'):
        bbox = bbox_from_npy(sys.argv[1])
    if sys.argv[1].endswith('.txt'):
        bbox = bbox_from_xml(sys.argv[1])
    else:
        bbox = bbox_from_xml(sys.argv[1])
else:
    bbox = bbox_from_xml()

bbox = np.array(bbox)
show_xml_bbox_space_distribute(bbox)


